#include <getopt.h>

#include "../ar8030_usb_test_tool/ar8030_com.h"
#include "CRC.h"
#include "bb_api.h"
#include "chk.hpp"
#include <fstream>
#include <inttypes.h>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include <semaphore.h>

#define WRITE_RETRY_TIMES 10

typedef struct {
    bb_dev_handle_t* pdev;
    int id;
    pthread_t pid;
    sem_t sem;
    uint32_t addr;
    unsigned int len;
    unsigned char* data;
} ar8030_upgrade_write_param_t;

typedef void *(*ThreadEntryPtrType)(void *param);

const int max_size = 0x300;
static int slot = -1;
static int err = 0;
static int written_size = 0;
static unsigned int total_len = 0;

static bb_host_t* phost;
static bb_dev_t** devs;

static int util_create_detach_thread(ThreadEntryPtrType entry, void *para, pthread_t *pid)
{
    pthread_t ThreadId;
    pthread_attr_t attr;

    pthread_attr_init(&attr);
    pthread_attr_setscope(&attr, PTHREAD_SCOPE_SYSTEM);
    pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
    if(pthread_create(&ThreadId, &attr, entry, para) == 0) {
        pthread_attr_destroy(&attr);
        if(pid)
            *pid = ThreadId;
        return 0;
    }
    pthread_attr_destroy(&attr);

    return -1;
}

int ar8030_ioctl(bb_dev_handle_t* pdev, uint32_t request, const void* in, uint16_t in_len, void* out, uint16_t out_len)
{
    int ret = 0;
    bb_remote_ioctl_in_t remote_in = {0};
    bb_remote_ioctl_out_t remote_out = {0};

    if (slot < 0) {
        ret = bb_ioctl(pdev, request, in, out);
    }
    else {
        memcpy(remote_in.data, in, in_len);
        remote_in.len = in_len;
        remote_in.slot = slot;
        remote_in.msg_id = request;

        ret = bb_ioctl(pdev, BB_REMOTE_IOCTL_REQ, &remote_in, &remote_out);
        if (out) {
            memcpy(out, remote_out.data, out_len);
        }
    }

    return ret;
}

int upgrade_runsys_get(bb_dev_handle_t* pdev, bb_runsys_t* runsys_idx)
{
    int ret = 0;
    bb_get_runsys_out_t runsys_out;

    if (!runsys_idx) {
        return -1;
    }

    ret = ar8030_ioctl(pdev, BB_GET_RUN_SYS, NULL, 0, &runsys_out, sizeof(bb_get_runsys_out_t));
    if (ret != 0) {
        printf("ar8030_ioctl get runsys failed.ret %d\n", ret);
    }

    *runsys_idx = runsys_out.runsys_id;

    return ret;
}

int ar8030_upgrade_write(bb_dev_handle_t* pdev, unsigned long long addr, unsigned short len, unsigned char* data)
{
    int retry_times = 0;
    int ret = 0;
    bb_set_hot_upgrade_write_in_t  in  = {0};
    bb_set_hot_upgrade_write_out_t out = {0};

    in.seq  = 5673;
    in.addr = (unsigned int)addr;
    in.len  = len;
    memcpy(in.data, data, len);

    do {
        ret = ar8030_ioctl(pdev, BB_SET_HOT_UPGRADE_WRITE, &in, sizeof(bb_set_hot_upgrade_write_in_t), &out, sizeof(bb_set_hot_upgrade_write_out_t));//bb_ioctl(pdev, BB_SET_HOT_UPGRADE_WRITE, &in, &out);
        if (ret || out.ret != 0) {
            retry_times++;
            if (retry_times > WRITE_RETRY_TIMES) {
                printf("BB_SET_HOT_UPGRADE_WRITE err! ret = %d out.ret = %d\n", ret, out.ret);
                return -1;
            }
        }
        else {
            break;
        }
    } while (1);

    return 0;
}

void *ar8030_upgrade_write_thread(void *param)
{
    unsigned int i          = 0;
    int          ret        = 0;
    unsigned int offset     = 0;
    ar8030_upgrade_write_param_t *write_param = (ar8030_upgrade_write_param_t*)param;

    if (!param) {
        return NULL;
    }

    if (!write_param->pdev || !write_param->data) {
        bb_dev_close(write_param->pdev);
        return NULL;
    }

    for (i = 0; i < (write_param->len / max_size); ++i) {
        // write 768
        if (err) {
            break;
        }

        offset = i * max_size;
        ret    = ar8030_upgrade_write(write_param->pdev, write_param->addr + offset, max_size, write_param->data + offset);
        if (ret < 0) {
            printf("ar8030_upgrade_write failed!\n");
            err = ret;
            break;
        }

        written_size += max_size;

        printf("update process %3.2f%%\r", (float)written_size * 100 / total_len);
        fflush(stdout);
    }

    if (write_param->len % max_size) {
        offset = i * max_size;
        ret  = ar8030_upgrade_write(write_param->pdev, write_param->addr + offset, write_param->len % max_size, write_param->data + offset);
        if (ret < 0) {
            printf("ar8030_upgrade_write failed!\n");
        }
        written_size += (write_param->len % max_size);
    }

    sem_post(&write_param->sem);

    return NULL;
}

int ar8030_upgrade_chk_crc(bb_dev_handle_t* pdev, uint32_t addr, unsigned int len, unsigned int target_crc)
{
    bb_set_hot_upgrade_crc32_in_t  in  = {};
    bb_set_hot_upgrade_crc32_out_t out = {};

    in.seq   = 3673;
    in.addr  = (unsigned int)addr;
    in.len   = len;
    in.crc32 = target_crc;
    int ret = ar8030_ioctl(pdev, BB_SET_HOT_UPGRADE_CRC32, &in, sizeof(bb_set_hot_upgrade_crc32_in_t), &out, sizeof(bb_set_hot_upgrade_crc32_out_t));//bb_ioctl(pdev, BB_SET_HOT_UPGRADE_WRITE, &in, &out);
    if (ret || out.ret != 0) {
        printf("BB_SET_HOT_UPGRADE_CRC32 err!\n");
        return -1;
    }

    return 0;
}

static int ar8030_data_transmit(uint32_t addr, unsigned int len, unsigned char* data, unsigned int* left_len)
{
    int          ret        = 0;
    int          i          = 0;
    unsigned int block_len  = 0;
    int          split      = 0;
    int          thread_num = 1;
    static ar8030_upgrade_write_param_t *write_param = NULL;

    if (!left_len) {
        printf("%s left_len is NULL!\n", __FUNCTION__);
        return -1;
    }

    /* Init the status */
    err = 0;

    if ((len / 0x1000) > BB_REMOTE_CMD_WAIT_MAX) {
        thread_num = BB_REMOTE_CMD_WAIT_MAX;
        block_len = (len / 0x1000) / thread_num * 0x1000;
    }
    else {
        thread_num = (len / 0x1000);
        block_len = 0x1000;
    }

    *left_len = (len - (block_len * thread_num));

    if (!write_param) {
        write_param = (ar8030_upgrade_write_param_t*)malloc(sizeof(ar8030_upgrade_write_param_t) * BB_REMOTE_CMD_WAIT_MAX);
        if (!write_param) {
            printf("%s write_param malloc failed!\n", __FUNCTION__);
            return -2;
        }
        memset(write_param, 0, sizeof(ar8030_upgrade_write_param_t) * BB_REMOTE_CMD_WAIT_MAX);

        for (i = 0;i < BB_REMOTE_CMD_WAIT_MAX;i++) {
            sem_init(&write_param[i].sem, 0, 0);
            write_param[i].addr = 0;
            write_param[i].data = 0;
            write_param[i].len  = 0;
            write_param[i].pdev = bb_dev_open(devs[0]);
            write_param[i].id   = 0;
        }
    }

    for (i = 0;i < thread_num;i++) {
        write_param[i].addr = addr + i*block_len;
        write_param[i].data = data + i*block_len;
        write_param[i].len  = block_len;
        write_param[i].id   = i;
        ret = util_create_detach_thread(ar8030_upgrade_write_thread, (void*)&write_param[i], &write_param[i].pid);
        if (ret) {
            printf("%s create thread failed! ret %d\n", __FUNCTION__, ret);
            err = -1;
            break;
        }
    }

    /* wait for threads */
    for (i--;i >=0;i--) {
        sem_wait(&write_param[i].sem);
    }

    return ret;
}

int ar8030_upgrade_partition(bb_dev_handle_t* pdev, uint32_t addr, unsigned int len, unsigned char* data)
{
    unsigned int o_crc      = 0;
    unsigned int i          = 0;
    unsigned int offset     = 0;
    unsigned int left_len   = len;
    unsigned int cur_len    = left_len;
    int          ret        = 0;

    // cal crc
    o_crc = CRC::Calculate(data, len, CRC::CRC_32());

    printf("%s %d addr %#" PRIx64 " len %d crc %08x\n", __func__, __LINE__, (uint64_t)addr, len, o_crc);

    written_size = 0;
    total_len = len;

    // split
    while (left_len > 0x1000) {
        ret = ar8030_data_transmit(addr + written_size, cur_len, data + written_size, &left_len);
        if (ret) {
            return -1;
        }
        cur_len = left_len;
    }

    for (; left_len > max_size; left_len-=max_size) {
        // write 768
        printf("update process %3.2f%%\r", (float)written_size * 100 / len);
        fflush(stdout);
        ret = ar8030_upgrade_write(pdev, addr + written_size, max_size, data + written_size);
        if (ret < 0) {
            printf("ar8030_upgrade_write failed!\n");
            return -1;
        }
        written_size += max_size;
    }

    if (left_len) {
        printf("update process %3.2f%%\r", (float)written_size * 100 / len);
        fflush(stdout);
        ret = ar8030_upgrade_write(pdev, addr + written_size, left_len, data + written_size);
        if (ret < 0) {
            printf("ar8030_upgrade_write failed!\n");
            return -1;
        }
    }
    printf("update process 100%%  \n");

#if 1
    // get crc
    ret = ar8030_upgrade_chk_crc(pdev, addr, len, o_crc);
    if (ret < 0) {
        printf("ar8030_upgrade_chk_crc failed!\n");
        return -1;
    }
#endif
    return 0;
}

int ar8030_upgrade_partition_byname(bb_dev_handle_t*     pdev,
                                    struct upgrade_hdr*  hdr,
                                    struct part_info*    partitions,
                                    struct segment_info* segments,
                                    unsigned char*       data,
                                    const char*          partname)
{
    int                i, ret = 0;
    int                segment_idx    = -1;
    unsigned long long partition_addr = 0;
    unsigned long long partition_len  = 0;

    for (i = 0; i < hdr->partitions; i++) {
        if (partitions[i].is_upgrade) {
            // printf("part %d, name %8s, length 0x%08llx, upgrade %d seg %d\n", i,
            //     partitions[i].name, partitions[i].length, partitions[i].is_upgrade, segment_idx);
            ++segment_idx;
        }

        if (0 == strcmp((char*)partitions[i].name, partname)) {
            break;
        }
    }

    // no not upgrade but return 0
    if (!partitions[i].is_upgrade)
        return 0;

    if (-1 == segment_idx || partitions[i].flash_offset != segments[segment_idx].flash_offset) {
        printf("partition %s info is not match! 0x%x 0x%x %d\n",
               partitions[i].name,
               (int)partitions[i].flash_offset,
               (int)segments[segment_idx].flash_offset,
               (int)segment_idx);
        return -1;
    }

    partition_addr = segments[segment_idx].flash_offset + GPT_FLASH_OFFSET;
    partition_len  = segments[segment_idx].size_decompress;
    data += segments[segment_idx].img_offset;
    ret = ar8030_upgrade_partition(pdev, partition_addr, partition_len, data);
    if (ret < 0) {
        printf("upgrade partition %s failed!\n", partname);
    }
    return ret;
}

int main(int argc, char** argv)
{
    int         ret             = 0;
    int         port            = BB_PORT_DEFAULT;
    const char* file_path       = "artosyn-upgrade-ar8030.img";
    const char* ip              = "127.0.0.1";
    uint32_t partition_addr     = 0;
    bb_runsys_t runsys_idx = BB_RUNSYS_MASTER;

    while (1) {
        int                  option_index   = 0;
        int                  c              = 0;
        static struct option long_options[] = {
            {"port",        required_argument, 0, 'p'},
            { "ip",         required_argument, 0, 'i'},
            { "file",       required_argument, 0, 'f'},
            { "slot",       required_argument, 0, 's'},
            { 0,            0,                 0, 0  },
        };

        c = getopt_long(argc, argv, "i:p:f:s:", long_options, &option_index);
        if (c == -1) {
            break;
        }

        switch (c) {
        case 'i':
            ip = optarg;
            break;
        case 'p':
            port = (int)strtoul(optarg, NULL, 10);
            break;
        case 'f':
            file_path = optarg;
            break;
        case 's':
            slot = (int)strtoul(optarg, NULL, 10);
            break;
        default:
            break;
        }
    }

    ret = bb_host_connect(&phost, ip, port);

    if (ret) {
        printf("connect failed = %d\n", ret);
        exit(-1);
    }

    int sz = bb_dev_getlist(phost, &devs);
    if (sz <= 0) {
        printf("dev cnt = 0\n");
        exit(-1);
    }

    bb_dev_handle_t* pdev = bb_dev_open(devs[0]);

    if (!pdev) {
        printf("can't open dev!!\n");
        exit(-1);
    }
    img8030 img(file_path);

    img.dump_upgrade_file();

    if (img.hdr->rom_size) {
        printf("update romcode datap = %p , len = %x\n", img.romcode, img.hdr->rom_size);
        ar8030_upgrade_partition(pdev, 0, img.hdr->rom_size, img.romcode);
    }

    ret = upgrade_runsys_get(pdev, &runsys_idx);
    if (ret != 0) {
        printf("get runsys failed! ret %d\n", ret);
        return -1;
    }

    printf("running sys is %d\n", runsys_idx);

    /** Upgrade the app that not using first to avoid the
     * case such as following:
     *  Sys is running in app1 since app0 has problem(not be
     * able to run). Then the power is powered off while
     * upgrade has been started and app1 in flash has been
     * overwritted. From now on, app0 and app1 will neither be
     * able to run.
    */
    if (runsys_idx == BB_RUNSYS_MASTER) {
        // write gpt table 1
        partition_addr = 0x0 + img.hdr->rom_size + GPT_FLASH_SIZE;
        ret                     = ar8030_upgrade_partition(pdev, partition_addr, GPT_FLASH_SIZE, img.gpt);
        if (ret < 0)
            return -1;
        // write partitions app1
        ret = ar8030_upgrade_partition_byname(pdev, img.hdr, img.partitions, img.segments, (unsigned char*)img.img, "app1");
        if (ret < 0)
            return -1;

        // write gpt table 0
        partition_addr = 0x0 + img.hdr->rom_size;
        ret            = ar8030_upgrade_partition(pdev, partition_addr, GPT_FLASH_SIZE, img.gpt);
        if (ret < 0)
            return -1;

        // write partitions app0
        ret = ar8030_upgrade_partition_byname(pdev, img.hdr, img.partitions, img.segments, (unsigned char*)img.img, "app0");
        if (ret < 0)
            return -1;
    }
    else {
        // write gpt table 0
        partition_addr = 0x0 + img.hdr->rom_size;
        ret            = ar8030_upgrade_partition(pdev, partition_addr, GPT_FLASH_SIZE, img.gpt);
        if (ret < 0)
            return -1;

        // write partitions app0
        ret = ar8030_upgrade_partition_byname(pdev, img.hdr, img.partitions, img.segments, (unsigned char*)img.img, "app0");
        if (ret < 0)
            return -1;

        // write gpt table 1
        partition_addr = 0x0 + img.hdr->rom_size + GPT_FLASH_SIZE;
        ret                     = ar8030_upgrade_partition(pdev, partition_addr, GPT_FLASH_SIZE, img.gpt);
        if (ret < 0)
            return -1;
        // write partitions app1
        ret = ar8030_upgrade_partition_byname(pdev, img.hdr, img.partitions, img.segments, (unsigned char*)img.img, "app1");
        if (ret < 0)
            return -1;
    }

    return 0;
}